package com.learning.datastruct.tree;

import java.util.Comparator;

/**
 * @author ADMIN
 * @Classname BinarySearchTree
 * @description 二叉搜索树
 * @date 2022/1/24 2:27
 */
@SuppressWarnings({"unchecked", "unused"})
public class BinarySearchTree<E> extends BinaryTree<E> {
    private Comparator<E> comparator = null;

    public BinarySearchTree() {
    }

    public BinarySearchTree(Comparator<E> comparator) {
        this.comparator = comparator;
    }

    public void add(E element) {
        elementNotNullCheck(element);
        Node<E> newNode = createNode(element, null);
        if (root == null) {
            root = newNode;
        } else {
            Node<E> node = root;
            Node<E> parent = null;
            int cmp = 0;
            while (node != null) {
                parent = node;
                cmp = compare(element, node.element);
                if (cmp > 0) {
                    node = node.right;
                } else if (cmp < 0) {
                    node = node.left;
                } else {
                    node.element = element;
                    return;
                }
            }
            newNode.parent = parent;
            if (cmp > 0) {
                parent.right = newNode;
            } else {
                parent.left = newNode;
            }
        }
        afterAdd(newNode);
        size++;
    }

    protected Node<E> createNode(E element, Node<E> parent) {
        return new Node<>(element, parent);
    }

    public void remove(E element) {
        remove(node(element));
    }

    public boolean contains(E element) {
        return node(element) != null;
    }

    protected void afterAdd(Node<E> newNode) {
    }

    protected void afterRemove(Node<E> newNode) {
    }

    private Node<E> node(E element) {
        Node<E> node = root;
        while (node != null) {
            int cmp = compare(element, node.element);
            if (cmp == 0) {
                return node;
            } else if (cmp > 0) {
                node = node.right;
            } else {
                node = node.left;
            }
        }
        return null;
    }

    private void remove(Node<E> node) {
        if (node == null) {
            return;
        }
        size--;
        if (node.hasTowChildren()) {
            Node<E> pre = precessNode(node);
            node.element = pre.element;
            node = pre;
        }
        Node<E> replacement = node.left != null ? node.left : node.right;
        if (replacement != null) {
            replacement.parent = node.parent;
            if (node.parent == null) {
                root = replacement;
            } else if (node == node.parent.left) {
                node.parent.left = replacement;
            } else {
                node.parent.right = replacement;
            }
            afterRemove(node);
        } else if (node.parent == null) {
            root = null;
        } else {
            if (node == node.parent.left) {
                node.parent.left = null;
            } else {
                node.parent.right = null;
            }
            afterRemove(node);
        }
    }


    private int compare(E e1, E e2) {
        if (this.comparator != null) {
            return this.comparator.compare(e1, e2);
        }
        return ((Comparable<E>) e1).compareTo(e2);
    }

    private void elementNotNullCheck(E element) {
        if (element == null) {
            throw new IllegalArgumentException("element must not be null");
        }
    }
}
